# backtester.py

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt


class Backtester:
    """
    执行向量化回测的核心引擎。
    """

    def __init__(self, data: pd.DataFrame, strategy, initial_capital: float, leverage: float = 1.0,
                 commission_rate: float = 0.0004):
        self.data = data
        self.strategy = strategy
        self.initial_capital = initial_capital
        self.leverage = leverage
        self.commission_rate = commission_rate
        self.results = None

    def run(self):
        """
        执行回测循环并计算绩效。
        这是一个简化的向量化回测实现。
        """
        print("生成策略信号...")
        signals = self.strategy.generate_signals(self.data)

        # 将信号与主数据合并
        self.data['signal'] = signals['signal']

        print("开始向量化回测...")

        # 计算持仓。信号的差异表示交易动作
        # signal从1到-1，diff为-2，表示从多头反手到空头
        # signal从0到1，diff为1，表示从空仓开多
        self.data['position'] = self.data.groupby('symbol')['signal'].transform(lambda x: x.shift(1)).fillna(0)

        # 计算每根K线的收益率
        self.data['returns'] = self.data.groupby('symbol')['close'].pct_change()

        # 策略收益率 = 持仓 * 下一根K线的收益率
        self.data['strategy_returns'] = self.data['position'] * self.data['returns']

        # --- 考虑交易成本 ---
        trades = self.data['position'].diff().abs()
        self.data['commission'] = trades * self.commission_rate

        # --- 考虑资金费用 ---
        # 资金费用在持仓期间产生，按持仓价值计算
        # 注意：这里的资金费率是针对上一周期的持仓
        self.data['funding_cost'] = self.data['position'] * self.data['fundingRate']

        # 计算净策略收益率
        self.data['net_strategy_returns'] = self.data['strategy_returns'] - self.data['commission'] - self.data[
            'funding_cost']

        # --- 计算权益曲线 ---
        # 假设资金平均分配到每个币种
        num_symbols = len(self.data.index.get_level_values('symbol').unique())

        # 计算每个币种的累计收益
        # 向量化回测中，我们通常计算每个时间点的累计收益率，然后乘以初始资本
        cumulative_returns = self.data.groupby(level='timestamp')['net_strategy_returns'].mean()

        # 杠杆效应
        cumulative_returns *= self.leverage

        equity_curve = (1 + cumulative_returns).cumprod() * self.initial_capital
        equity_curve.name = 'equity'

        # 将权益曲线存入结果
        self.results = pd.DataFrame(equity_curve)

        print("回测完成。")
        return self._calculate_performance()

    def _calculate_performance(self):
        """计算并返回关键绩效指标。"""
        if self.results is None:
            raise Exception("请先运行回测！")

        equity = self.results['equity']

        # 总收益率
        total_return = (equity.iloc[-1] / equity.iloc[0] - 1) * 100

        # 最大回撤
        rolling_max = equity.cummax()
        drawdown = (equity - rolling_max) / rolling_max
        max_drawdown = drawdown.min() * 100

        # 夏普比率 (假设无风险利率为0，每日数据)
        daily_returns = equity.pct_change().dropna()
        # 年化夏普比率
        sharpe_ratio = (daily_returns.mean() / daily_returns.std()) * np.sqrt(365)  # 假设数据是每日的

        performance = {
            "初始资本": f"${self.initial_capital:,.2f}",
            "期末权益": f"${equity.iloc[-1]:,.2f}",
            "总收益率 (%)": f"{total_return:.2f}",
            "最大回撤 (%)": f"{max_drawdown:.2f}",
            "年化夏普比率 (无风险利率=0)": f"{sharpe_ratio:.2f}"
        }
        return performance

    def plot_equity_curve(self):
        """绘制资金权益曲线。"""
        if self.results is None:
            raise Exception("请先运行回测！")

        plt.figure(figsize=(15, 7))
        self.results['equity'].plot()
        plt.title('资金权益曲线 (Equity Curve)')
        plt.xlabel('日期')
        plt.ylabel('权益价值 ($)')
        plt.grid(True)
        plt.show()